load("//bazel:api.bzl", "mojo_library", "requirement")
load(":custom_op_example.bzl", "custom_op_example_py_binary")

package(default_visibility = ["//visibility:private"])

filegroup(
    name = "kernel_sources",
    srcs = glob(["operations/*.mojo"]),
)

custom_op_example_py_binary(
    name = "addition",
    srcs = ["addition.py"],
    extra_data = [
        ":operations",
    ],
    extra_deps = [
        "//max/python/max/torch",
        requirement("torch"),
    ],
)

custom_op_example_py_binary(
    name = "grayscale",
    srcs = ["grayscale.py"],
    extra_data = [
        ":operations",
    ],
    extra_deps = [
        "//max/python/max/torch",
        requirement("torch"),
        requirement("pillow"),
    ],
)

custom_op_example_py_binary(
    name = "whisper",
    srcs = ["whisper.py"],
    create_test = False,
    extra_data = [
        ":operations",
    ],
    extra_deps = [
        "//max/python/max/torch",
        requirement("torch"),
        requirement("transformers"),
        requirement("datasets"),
        requirement("librosa"),
    ],
)

custom_op_example_py_binary(
    name = "graph",
    srcs = ["graph.py"],
    extra_deps = [
        "//max/python/max/torch",
        requirement("torch"),
    ],
)

mojo_library(
    name = "operations",
    srcs = [":kernel_sources"],
    validate_missing_docs = False,
    deps = [
        "//max:compiler",
        "//max:layout",
        "//max:tensor",
        "@mojo//:stdlib",
    ],
)
